package com.ivy.springai.controller;

import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;

@RestController
@RequestMapping("/chat/stream")
public class StreamingChatClientController {

    @Autowired
    private StreamingChatClient streamingChatClient;//流式处理聊天客户端

    @GetMapping(value = "/generate", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        System.out.println("调用流式输出："+message);
        Prompt prompt = new Prompt(new UserMessage(message), OpenAiChatOptions.builder()
                .build());
        return streamingChatClient.stream(prompt);
    }

    @GetMapping(value = "/generate2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<OpenAiApi.ChatCompletionChunk> generate2(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        System.out.println("调用流式输出："+message);
        OpenAiApi openAiApi = new OpenAiApi(System.getenv("OPENAI_API_KEY"));
        OpenAiApi.ChatCompletionMessage chatCompletionMessage = new OpenAiApi.ChatCompletionMessage("Hello world", OpenAiApi.ChatCompletionMessage.Role.USER);
        // Sync request
        ResponseEntity<OpenAiApi.ChatCompletion> response = openAiApi.chatCompletionEntity(
                new OpenAiApi.ChatCompletionRequest(List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8f, false));

        // Streaming request
        Flux<OpenAiApi.ChatCompletionChunk> streamResponse = openAiApi.chatCompletionStream(
                new OpenAiApi.ChatCompletionRequest(List.of(chatCompletionMessage), "gpt-3.5-turbo", 0.8f, true));
        return streamResponse;
    }
}
